package com.xinan.ollama.api.controller;

import com.baomidou.mybatisplus.extension.api.R;
import com.xinan.ollama.api.component.OllamaComponent;
import com.xinan.ollama.api.vo.req.OllamaReq;
import io.swagger.annotations.ApiOperation;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import javax.annotation.Resource;
import java.util.Map;

/**
 * Ollama接口
 *
 * @author yuanz
 * @date 2025-02-21
 */
@Slf4j
@RestController
@RequestMapping("/api")
public class OllamaController {

    @Resource
    private OllamaComponent ollamaComponent;

    @PostMapping("/generate")
    @ApiOperation(value = "生成补全", tags = {"baseChat"})
    public R<Object> generate(@RequestBody OllamaReq req) {
        return ollamaComponent.generate(req);
    }

    @PostMapping("/chat")
    @ApiOperation(value = "生成聊天补全", tags = {"baseChat"})
    public R<Object> chat(@RequestBody OllamaReq req) {
        return ollamaComponent.chat(req);
    }

    @PostMapping("/getEmbeddingByStr")
    @ApiOperation(value = "获取字符串的嵌入向量", tags = {"RAG"})
    public R<Object> getEmbeddingByStr(@RequestBody OllamaReq req) {
        return ollamaComponent.getEmbeddingByStr(req);
    }

    @PostMapping("/saveStrEmbeddingToDB")
    @ApiOperation(value = "获取并保存字符串的嵌入向量", tags = {"RAG"})
    public R<Object> saveStrEmbeddingData(@RequestBody OllamaReq req) {
        return ollamaComponent.saveStrEmbeddingData(req);
    }

    @PostMapping("/chatOfRag")
    @ApiOperation(value = "生成聊天补全（RAG）", tags = {"RAG"})
    public R<Object> chatOfRag(@RequestBody OllamaReq req) {
        return ollamaComponent.chatOfRag(req);
    }

    @PostMapping("/getEventAndOrganizeNameByInformationName")
    @ApiOperation(value = "通过情报名称获取事件及组织名称", tags = {"KeyPerson"})
    public R<Object> getEventAndOrganizeNameByInformationName(@RequestBody OllamaReq req) {
        return ollamaComponent.getEventAndOrganizeNameByInformationName(req);
    }

    @PostMapping("/chatForKeyPerson")
    @ApiOperation(value = "重点人员数据提取", tags = {"KeyPerson"})
    public R<Object> chatForKeyPerson(@RequestBody OllamaReq req) {
        return ollamaComponent.chatForKeyPerson(req);
    }

    @PostMapping("/chatForAnalysisReport")
    @ApiOperation(value = "重点人员分析报告", tags = {"KeyPerson"})
    public R<Object> chatForAnalysisReport(@RequestBody OllamaReq req) {
        return ollamaComponent.chatForAnalysisReport(req);
    }

    @PostMapping("/chatForPersonLevel")
    @ApiOperation(value = "重点人员等级判定", tags = {"KeyPerson"})
    public R<Object> chatForPersonLevel(@RequestBody OllamaReq req) {
        return ollamaComponent.chatForPersonLevel(req);
    }

    @PostMapping("/generateInformationContentByData")
    @ApiOperation(value = "通过数据生成情报内容", tags = {"KeyPerson"})
    public R<Object> generateInformationContentByData(@RequestBody OllamaReq req) {
        return ollamaComponent.generateInformationContentByData(req);
    }

    @PostMapping("/generateSubjectAnalysisReportByParam")
    @ApiOperation(value = "通过参数生成专题分析报告内容", tags = {"KeyPerson"})
    public R<Object> generateSubjectAnalysisReportByParam(@RequestBody Map<String, Object> req) {
        return ollamaComponent.generateSubjectAnalysisReportByParam(req);
    }

}
